"""
本文件是构建faiss数据库 + 测试faiss数据库的效果
"""

from main import get_dataloader
import faiss
from CNN_NET import *
import json

# 统计一个数组中的出现最多的索引(得分法)
def find_most(list):
    if len(list) == 0:
        return -1
    score = 0
    top = list[0] #得分主
    for e in list:
        if top == e:
            score = score+1
        else:
            score = score-1
            if score < 0:
                score = 1
                top = e
    return top

# 创建faiss数据库
def creat_faiss(net_path=r"D:\Code\CNN识别菌种\Model\E20_2024-06-04_20.37.27.pth"):
    # 1. 加载数据集和测试集
    TrainDataLoader, TestDataLoader, net_map, KindIndexMap = get_dataloader()

    # 2. 加载网络
    net = MyNet(len(net_map), 51, 47)
    net = torch.load(net_path)
    net.eval()

    # 3. 创建faiss数据库index
    features_list = []  # 保存特征的列表或数组
    id_to_kind = [] #索引对应菌种
    feature_dim = 50

    for (x,y) in TrainDataLoader:
        feature = net.get_features(x)
        features_list.append(feature.detach().numpy())
        id_to_kind.append(y[0])
    all_features = np.vstack([features for features in features_list])

    index = faiss.IndexFlatL2(feature_dim)  # 使用L2距离和没有量化的索引
    index.add(all_features)    # 添加所有特征到faiss索引

    return net,index,id_to_kind,TestDataLoader

# faiss测试,index是faiss数据库,id_to_kind是数据库的菌种对应索引
def test_faiss(net, index, id_to_kind, TestDataLoader):
    k = 3 # 搜索最近的3个种
    n_correct = 0

    for (x, y) in TestDataLoader:
        print("-----------------------------------------------------------------")
        feature = net.get_features(x)
        D, I = index.search(feature.detach().numpy(), k)
        ans = []
        for id in I[0]:
            ans.append(id_to_kind[id])
        predict = find_most(ans)
        print("FAISS数据库相近的物种为：", ans)
        print("FAISS数据库预测：",predict)
        print("实际物种为：", y)

        if predict == y[0]:
            n_correct = n_correct + 1

    print("FAISS数据库，准确个数：",n_correct,"  共计：",TestDataLoader.__len__(),"准确率为：",n_correct/TestDataLoader.__len__())
    return n_correct/TestDataLoader.__len__()

# 测试程序
def test_demo():
    acc_list = []
    for _ in range(20):
        net, index, id_to_kind, TestDataLoader = creat_faiss()
        acc = test_faiss(net, index, id_to_kind, TestDataLoader)
        acc_list.append(acc)

    acc_list.append(sum(acc_list)/len(acc_list))

    #写入准确率文件
    with open('RunningData\\faiss.json', 'w') as f:
        json.dump(acc_list, f)


test_demo()
